#!/bin/bash

RESERVEDRAW=68
DEFAULT_LINUX_PARTITION="4P=200M:2G:5G:-1"
diskpartition=$DEFAULT_LINUX_PARTITION

usage ()
{
cat <<EOF
Usage:  installer [ -r <rootfs> ] [ -d <device> ] [ -t <types> ]
OPTION:
 -r, --rootfs       The first distro rootfs image by default to be programmed into target storage drive
 -d, --device       Device name of the target NVME/SATA storage drive in Linux
 -t, --types         uefi/uboot as bootloader
Examples:
  ./installer  -r rootfs-ubuntu-desktop.tar.gz -d sata (or nvme0)
  ./installer  -r phytium-image-lxqt-<machine>.tar.gz -d sata (or nvme0) 

EOF
    exit
}

search_dev() {
    if [ $1 = sata ]; then
        devlist=`ls /dev/sd?`	
        for dev in $devlist; do
            sdx=`udevadm info -q all -n $dev | grep ID_BUS`
            disksize=$[ `sudo fdisk -ls $dev` / 1000000 ]
            if [ $1 = sata ]; then
                if [ ${sdx:10:3} = ata ]; then
                    devsata=${dev:5:3}
                    echo "SATA disk: $devsata $disksize GB"
                    break
                fi
	    fi
        done
    else
	dev=`ls /dev/nvme0n1`
	disksize=$[ `sudo fdisk -ls $dev` / 1000000 ]
	devnvme=${dev:5:7}
	echo "NVME disk: $devnvme $disksize GB"
    fi
}

search_usb() {
    devlist=`ls /dev/sd?`
    for dev in $devlist; do
        sdx=`udevadm info -q all -n $dev | grep ID_BUS`
        disksize=$[ `sudo fdisk -ls $dev` / 1000000 ]
        if [ $1 = usb ] ; then
            if [ ${sdx:10:3} = usb ]; then
                devusb=${dev:5:3}
                echo "USB disk: $devusb $disksize GB"
                break
            fi
        elif [ $1 = sata ]; then
            if [ ${sdx:10:3} = ata ]; then
                devsata=${dev:5:3}
                echo "SATA disk: $devsata $disksize GB"
                break
            fi
        fi
    done
}


check_partition_number() {
    [ -z "$diskpartition" ] && print_e "Please specify partition settings" && exit

    if ! [[ $diskpartition =~ ^[2-9]+[P=]+[0-9]+[G,M]* ]] || \
        echo $diskpartition | tr -d '[0-9] [:,-]' | grep  -E '[^P,M,G]'; then
        print_e "Error: the number of partitions should be >= 2 and the unit of size is G or M"
        print_w "Correct example: '-p 5P=200M:2G:5G:6G:-1', '-1' indicates the rest space of target disk for the last partition as rootfs"
        exit
    fi
    pnum=`echo $diskpartition | cut -d= -f1`
    pnum=${pnum%?}
    partlistnum=$(echo $diskpartition | awk -F":" '{print NF-1}')
    partlistnum=$[ $partlistnum + 1 ]
    [ $pnum != $partlistnum ] && print_e "ERROR: You set $pnum partitions, but listed $partlistnum partitions" && exit || true
    [ $pnum -ge 5 ] && totalparts=$[ $pnum + 1 ] || totalparts=$pnum
}

umount_device() {
    mntlist=$(cat /proc/mounts | grep /dev/${devpartname} | cut -d' ' -f1 | tr '\n' ' ')
    for mntp in $mntlist; do
        sudo fuser -k $mntp || true
        sudo umount $mntp || true
    done
}


mount_device() {
    for ((i=1; i<=$pnum; i++)); do
        [ $pnum -ge 5 -a $i = 4 ] && continue
        sudo mkdir -p $mntdir/${devpartname}$i
        if ! mount | grep -q ${devpartname}$i; then
            sudo mount /dev/${devpartname}$i $mntdir/${devpartname}$i 1>/dev/null 2>&1 || true
        fi
    done
}


partition_format_disk() {
    if mount | grep -q "$sysdev on / type"; then
        print_w "The target /dev/$devname can't be the running device in which $sysdev is running as rootfs"
        print_w "You can choose a non-running device or do the installation in Linux environment" && exit
    fi
    print_n "Partitioning /dev/$devname ..."

    optimal_io_size=$(sudo cat /sys/block/$devname/queue/optimal_io_size)
    minimum_io_size=$(sudo cat /sys/block/$devname/queue/minimum_io_size)
    [ "$optimal_io_size" = "0" ] && aligntype=minimal || aligntype=optimal
    [ "$enable_efi" = "y" ] && parttable=gpt || parttable=msdos

    umount_device
    sudo rm -rf /mnt/${devpartname}*/*

    sudo parted -a $aligntype -s /dev/$devname mklabel $parttable

    for ((i=1; i<=$pnum; i++)); do
        eval n_p$i=`echo $1 | cut -d: -f$i`
        [ $i = 1 ] && n_p1=`echo $n_p1 | cut -d= -f2`
        n_p=`eval echo '${n_p'"$i"'}'`
        [ ${n_p: -1} = G ] && n_p=${n_p%?} && n_p=$[ $n_p * 1024 ]
        [ ${n_p: -1} = M ] && n_p=${n_p%?}
        [ ${n_p: -2} = -1 ] && eval e_p$i=100%
        [ $i = 1 ] && s_p1=$RESERVEDRAW && e_p1=$[ $s_p1 + $n_p ] && umount_device && \
        sudo parted -a $aligntype -s /dev/$devname mkpart primary ${s_p1}MiB ${e_p1}MiB && continue

        n=$[ i-1 ] && a_e_p=`eval echo '${e_p'"$n"'}'`
        eval s_p$i=$[ $a_e_p + 1 ] && s_p=`eval echo '${s_p'"$i"'}'`
        [ ${n_p: -2} = -1 ] && eval e_p$i=100% || eval e_p$i=$[ $s_p + $n_p ]
        e_p=`eval echo '${e_p'"$i"'}'`

        if [ $pnum -le 4 ]; then
            parttype=primary
        elif [ $i -le 3 ]; then
            parttype=primary
        elif [ $i -eq 4 ]; then
            parttype=extended
            end=100%
            umount_device
            sudo parted -a $aligntype -s /dev/$devname mkpart $parttype ${s_p}MiB $end
            s_p=$[ $s_p + 1 ]
            parttype=logical
            extendedpart=4
        else
            parttype=logical
        fi
        [ $e_p != 100% ] && end=${e_p}MiB || end=${e_p}
        umount_device
        sudo parted -a $aligntype -s /dev/$devname mkpart $parttype ${s_p}MiB $end
    done

    print_n "Formatting partitions ..."
    for ((i=1; i<=$totalparts; i++)); do
        if [ $i = "$bootpartnum" ]; then
            umount_device
            sudo mkfs.ext4 -F -q -b 4096 -L boot -O ^64bit$metadataopt $bootdev || true
        elif [ $i = "$efipartnum" ]; then
            if [ "$enable_efi" = "y" ]; then
                # in case of UEFI as Bootloader
                umount_device
                sudo mkfs.vfat -n EFI $efidev || true
            else
                # for misc metadata or other uses in case of U-Boot as Bootloader
                umount_device
                sudo mkfs.ext4 -F -q -b 4096 -L misc $efidev || true
            fi
        else
            if [ $i = "$syspartnum" ]; then
                label=system
            elif [ $i = "$backuppartnum" ]; then
                label=backup
            elif [ $i = "$userpartnum" ]; then
                label=userdata
            else
                label=data$i
            fi
            umount_device
            [ $i != "$extendedpart" ] && sudo mkfs.ext4 -F -q -O ^huge_file,^64bit$metadataopt -b 4096 -L $label /dev/${devpartname}$i || true
        fi
    done
    mount_device
    sudo parted -s /dev/$devname print
    print_d "partition and format /dev/$devname"
}

flex_install_distro() {
     
    search_usb usb
    devlist=`ls /dev/$devusb?`
    for dev in $devlist; do
        partuuid_sys=`lsblk -l --output UUID $dev | grep -`
        if [ ${partuuid_sys} = "5b8015f3-84e8-4445-9ed0-5c478cf88c87" ] ; then
            devusb=${dev:5:4}
            bootpart_n=/run/media/${devusb}
        fi
    done

    # install boot partition
    if [ -d $bootpart_n -o -f $bootpart_n ]; then
        sudo rm -rf $bootpartition/*
        echo Installing $bootpart_n to $bootdev, waiting ...
        sudo cp -rfp $bootpart_n/* $bootpartition
        print_d "Install $bootpart_n in $bootdev"
    fi

    # install system rootfs
    if [ -d $rootfs_n -o -f $rootfs_n ]; then
        echo Installing $rootfs_n to $sysdev, waiting ...
        if mount | grep -q "$sysdev on / type"; then
            print_w "Can't install rootfs to the running root filesystem" && exit
        fi
        sudo tar xf $rootfs_n -C $syspartition
        print_d "Install $rootfs_n in $sysdev"
    fi
    if [ "$enable_efi" = y ]; then
        # configure grub.cfg for UEFI
        if [ -f $bootpartition/EFI/BOOT/grub.cfg ]; then
	    partuuid_sys=`lsblk -l --output PARTUUID $sysdev | grep -`
            grep -l 'console=ttyAMA1,115200' $bootpartition/EFI/BOOT/grub.cfg | \
            sudo xargs sed -i "s#5b8015f3-84e8-4445-9ed0-5c478cf88c88#$partuuid_sys#g"
        fi
        if ! mount | grep ${devpartname}$efipartnum; then
           sudo mount $efidev $efipartition
        fi
        sudo mkdir -p $efipartition/EFI/BOOT
        if [ -f $bootpartition/EFI/BOOT/grub.cfg ]; then
            sudo cp $bootpartition/EFI/BOOT/grub.cfg $efipartition/EFI/BOOT/grub.cfg
        fi
        if [ -f $bootpartition/EFI/BOOT/bootaa64.efi ]; then
            sudo cp $bootpartition/EFI/BOOT/bootaa64.efi $efipartition/EFI/BOOT/
        fi
	if [ -f $bootpartition/EFI/BOOT/startup.nsh ]; then
            sudo cp $bootpartition/EFI/BOOT/startup.nsh $efipartition/EFI/BOOT/
        fi
	if [ -f $bootpartition/Image ]; then
            sudo cp $bootpartition/Image $efipartition
        fi
	if [ -f $bootpartition/initrd.img ]; then
            sudo cp $bootpartition/initrd.img $efipartition
        fi
    fi

}

print_e() {
    echo -e "${RED}$1 ${NC}"
}

print_n() {
    echo -e "${green}$1 ${NC}"
}

print_w() {
    echo -e "${YELLOW}$1 ${NC}"
}

print_d() {
    echo -e "${GREEN}$1     [Done] ${NC}"
}

RED='\e[1;31m'
GREEN='\e[1;32m'
green='\e[0;32m'
YELLOW='\e[1;33m'
NC='\e[0m'

[ -z "$1" -o "$1" = "-h" -o "$1" = "--help" ] && usage && exit

ARGS=$(getopt -a -o r:d:t:h \
       -l rootfs:,device:,types:,help -- "$@")

[ $? -ne 0 ] && usage
eval set -- "${ARGS}"

while true
do
        case "$1" in
        -r|--rootfs)
                rootfs=$2; shift;;
        -d|--device)
                mediadevice=$2; shift;;
	-t|--types)
                types=$2; shift;;
        -h|--help)
                usage;;
        --)
                shift; break;;
        esac
shift
done


if [ -n "$rootfs" ]; then
    search_usb usb
    devlist=`ls /dev/$devusb?`
    for dev in $devlist; do
        partuuid_sys=`lsblk -l --output UUID $dev | grep -`
        if [ ${partuuid_sys} = "5b8015f3-84e8-4445-9ed0-5c478cf88c89" ] ; then
            devusb=${dev:5:4}
            rootfs_n=/run/media/${devusb}/$rootfs
        fi
    done
fi

[ ! -f "$rootfs_n" ] && \
    print_w "Not found valid rootfs installed on usb, please specify '-r <rootfs>'" && exit

if [ "$mediadevice" = "sata" ] ; then
    search_dev sata
    devname=$devsata
fi

if [ "$mediadevice" = "nvme0" ] ; then
    search_dev nvme0
    devname=$devnvme
fi

[ ${devname:0:7} = nvme0n1 -o ${devname:0:4} = loop ] && devpartname=${devname}p || devpartname=$devname

enable_efi=y

check_partition_number

if [ $pnum -le 4 ]; then
        efipartnum=1; bootpartnum=2; backuppartnum=$[ $pnum - 1 ]; syspartnum=$pnum
    else
        efipartnum=1; bootpartnum=2; backuppartnum=$pnum; syspartnum=$[ $pnum + 1 ]
fi

mntdir=/mnt
efidev=/dev/${devpartname}$efipartnum
bootdev=/dev/${devpartname}$bootpartnum
sysdev=/dev/${devpartname}$syspartnum
backupdev=/dev/${devpartname}$backuppartnum
efipartition=$mntdir/${devpartname}${efipartnum}
bootpartition=$mntdir/${devpartname}${bootpartnum}
syspartition=$mntdir/${devpartname}${syspartnum}
backuppartition=$mntdir/${devpartname}${backuppartnum}

for i in $bootpartition $syspartition $backuppartition $efipartition ; do
    sudo mkdir -p $i
done

partition_format_disk $diskpartition
flex_install_distro
